"""
Claude LangChain Adapter - Makes Claude API compatible with LangChain
"""

import logging
import aiohttp
from typing import AsyncGenerator, List, Optional, Dict, Any, Union
from langchain.schema import BaseMessage, HumanMessage, SystemMessage, AIMessage
from langchain.schema.messages import AIMessageChunk
from langchain.callbacks.manager import AsyncCallbackManagerForLLMRun
from langchain_core.runnables import Runnable
from langchain_core.prompt_values import ChatPromptValue
from langchain_core.messages import BaseMessage as CoreBaseMessage

logger = logging.getLogger(__name__)


class ClaudeLangChainAdapter(Runnable):
    """Adapter to make Claude API work with LangChain chains"""
    
    def __init__(
        self,
        model_name: str,
        api_key: str,
        api_base: str,
        temperature: float = 0.7,
        max_tokens: int = 2000
    ):
        self.model_name = model_name
        self.api_key = api_key
        self.api_base = api_base.rstrip('/')
        self.temperature = temperature
        self.max_tokens = max_tokens
        
        # Ensure correct endpoint for Claude
        if not self.api_base.endswith('/v1/messages'):
            if '/v1' in self.api_base:
                self.api_base = self.api_base.replace('/v1', '/v1/messages')
            else:
                self.api_base = self.api_base + '/v1/messages'
    
    async def astream(
        self,
        input: Union[Dict, List[BaseMessage], ChatPromptValue],
        config: Optional[Dict] = None,
        **kwargs
    ) -> AsyncGenerator[Any, None]:
        """Stream responses from Claude API - Runnable interface"""
        
        logger.info(f"[ClaudeLangChainAdapter] astream called with input type: {type(input)}")
        
        # Extract messages from various input formats
        if isinstance(input, dict):
            # Handle dict input (from prompt template)
            messages = input.get('messages', [])
            logger.info(f"[ClaudeLangChainAdapter] Extracted {len(messages)} messages from dict")
        elif isinstance(input, ChatPromptValue):
            # Handle ChatPromptValue
            messages = input.messages
            logger.info(f"[ClaudeLangChainAdapter] Extracted {len(messages)} messages from ChatPromptValue")
        elif isinstance(input, list):
            # Direct list of messages
            messages = input
            logger.info(f"[ClaudeLangChainAdapter] Using list directly with {len(messages)} messages")
        else:
            # Try to extract messages attribute
            messages = getattr(input, 'messages', [input])
            logger.info(f"[ClaudeLangChainAdapter] Extracted {len(messages)} messages from attribute")
        
        # Convert LangChain messages to Claude format
        system_message = None
        claude_messages = []
        
        for msg in messages:
            logger.debug(f"[ClaudeLangChainAdapter] Processing message type: {type(msg)}")
            if hasattr(msg, 'type') and hasattr(msg, 'content'):
                # Handle CoreBaseMessage
                msg_type = msg.type
                msg_content = msg.content
                logger.debug(f"[ClaudeLangChainAdapter] CoreBaseMessage - type: {msg_type}, content length: {len(msg_content) if msg_content else 0}")
                
                if msg_type == 'system':
                    # Claude API expects system message separately
                    system_message = msg_content
                elif msg_type == 'human':
                    claude_messages.append({"role": "user", "content": msg_content})
                elif msg_type == 'ai':
                    claude_messages.append({"role": "assistant", "content": msg_content})
                else:
                    claude_messages.append({"role": "user", "content": msg_content})
            elif isinstance(msg, (SystemMessage, HumanMessage, AIMessage)):
                # Handle legacy message types
                if isinstance(msg, SystemMessage):
                    system_message = msg.content
                elif isinstance(msg, HumanMessage):
                    claude_messages.append({"role": "user", "content": msg.content})
                elif isinstance(msg, AIMessage):
                    claude_messages.append({"role": "assistant", "content": msg.content})
            else:
                # Default to user message
                claude_messages.append({"role": "user", "content": str(msg)})
        
        # Claude API payload
        payload = {
            "model": self.model_name,
            "messages": claude_messages,
            "max_tokens": self.max_tokens,
            "temperature": self.temperature,
            "stream": True
        }
        
        # Add system message if present
        if system_message:
            payload["system"] = system_message
        
        # Log the actual messages being sent
        logger.info(f"[ClaudeLangChainAdapter] Claude messages to send:")
        for i, msg in enumerate(claude_messages):
            role = msg.get('role', 'unknown')
            content = msg.get('content', '')
            logger.info(f"  Message {i+1}: role={role}, content_length={len(content)}")
            if content:
                logger.debug(f"  Content preview: {content[:200]}...")
        
        headers = {
            "x-api-key": self.api_key,  # Claude uses x-api-key header
            "anthropic-version": "2023-06-01",  # Required version header
            "Content-Type": "application/json"
        }
        
        logger.info(f"[ClaudeLangChainAdapter] Calling Claude API at {self.api_base}")
        logger.info(f"[ClaudeLangChainAdapter] Model: {self.model_name}, Messages count: {len(claude_messages)}")
        logger.info(f"[ClaudeLangChainAdapter] API Key prefix: {self.api_key[:10] if self.api_key else 'None'}...")
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    self.api_base,
                    headers=headers,
                    json=payload,
                    timeout=aiohttp.ClientTimeout(total=60)
                ) as response:
                    if response.status != 200:
                        error_text = await response.text()
                        logger.error(f"Claude API error {response.status}: {error_text}")
                        yield f"Error: Claude API returned {response.status}: {error_text}"
                        return
                    
                    logger.info(f"[ClaudeLangChainAdapter] Claude API response status: {response.status}")
                    logger.info(f"[ClaudeLangChainAdapter] Response headers: {dict(response.headers)}")
                    
                    # Handle streaming response
                    buffer = b''
                    chunk_count = 0
                    total_content = []
                    
                    async for chunk in response.content.iter_any():
                        buffer += chunk
                        lines = buffer.split(b'\n')
                        buffer = lines[-1]  # Keep incomplete line in buffer
                        
                        for line in lines[:-1]:
                            line_text = line.decode('utf-8').strip()
                            if line_text.startswith('data: '):
                                data_str = line_text[6:]
                                if data_str == '[DONE]':
                                    logger.info(f"[ClaudeLangChainAdapter] Received [DONE] signal after {chunk_count} chunks")
                                    logger.info(f"[ClaudeLangChainAdapter] Total content generated: {len(''.join(total_content))} characters")
                                    return
                                
                                try:
                                    import json
                                    data = json.loads(data_str)
                                    
                                    chunk_type = data.get('type')
                                    logger.debug(f"[ClaudeLangChainAdapter] Received chunk type: {chunk_type}, data keys: {list(data.keys())}")
                                    
                                    # Handle Claude streaming format
                                    if chunk_type == 'content_block_delta':
                                        delta = data.get('delta', {})
                                        if 'text' in delta:
                                            chunk_count += 1
                                            text = delta['text']
                                            total_content.append(text)
                                            logger.debug(f"[ClaudeLangChainAdapter] Yielding text chunk {chunk_count}: '{text[:50]}'...")
                                            # Create an object that mimics AIMessage chunk structure
                                            yield AIMessageChunk(content=text)
                                    elif chunk_type == 'message_delta':
                                        # Sometimes Claude uses this format
                                        delta = data.get('delta', {})
                                        if 'content' in delta:
                                            chunk_count += 1
                                            content = delta['content']
                                            total_content.append(content)
                                            logger.debug(f"[ClaudeLangChainAdapter] Yielding message chunk {chunk_count}: '{content[:50]}'...")
                                            yield AIMessageChunk(content=content)
                                    elif chunk_type == 'message_start':
                                        logger.debug(f"[ClaudeLangChainAdapter] Message started with data: {data}")
                                    elif chunk_type == 'content_block_start':
                                        logger.debug(f"[ClaudeLangChainAdapter] Content block started with data: {data}")
                                    elif chunk_type == 'content_block_stop':
                                        logger.debug(f"[ClaudeLangChainAdapter] Content block stopped")
                                    elif chunk_type == 'message_stop':
                                        logger.debug(f"[ClaudeLangChainAdapter] Message stopped")
                                    else:
                                        logger.debug(f"[ClaudeLangChainAdapter] Unknown chunk type: {chunk_type}, data: {data}")
                                except json.JSONDecodeError as e:
                                    logger.warning(f"[ClaudeLangChainAdapter] Failed to parse Claude stream data: {data_str[:100]}, error: {e}")
                                    continue
                    
                    # Process remaining buffer
                    if buffer:
                        logger.debug(f"[ClaudeLangChainAdapter] Processing remaining buffer: {buffer[:100]}")
                    
                    logger.info(f"[ClaudeLangChainAdapter] Stream completed, total chunks yielded: {chunk_count}")
                    if total_content:
                        logger.info(f"[ClaudeLangChainAdapter] Total content preview: {(''.join(total_content))[:200]}...")
                                    
        except Exception as e:
            logger.error(f"Claude API call failed: {e}")
            yield f"Error: {str(e)}"
    
    async def ainvoke(self, input, config=None, **kwargs):
        """Non-streaming invoke (collects all chunks)"""
        chunks = []
        async for chunk in self.astream(input, config, **kwargs):
            chunks.append(chunk)
        return ''.join(chunks)
    
    def invoke(self, input, config=None, **kwargs):
        """Synchronous invoke - not implemented for async adapter"""
        raise NotImplementedError("ClaudeLangChainAdapter only supports async operations")